#include "pcap.h"
#include "packet32.h"
#include "ntddndis.h"

#include "struct.h"
#include "init.h"
#include "checksum.h"
#include "regroup.h"
#include "arp.h"
#include "hash.h"

int opened = 0;
table_entry *port2ent[65536];
table_entry *ent2port[9127];
subnet_entry *subn[233];

void subnet_put(ip_address ip, mac_address mac) {
	subnet_entry *i, *ne;

	if (subn[hash2(ip)] != 0) {
		i = subn[hash2(ip)];

		while (i) {
			if (ip.byte1 == i->ip.byte1 && ip.byte2 == i->ip.byte2 && ip.byte3 == i->ip.byte3 && ip.byte4 == i->ip.byte4) return;
			i = i->next;
		}
	}

	ne = new subnet_entry;
	ne->ip = ip;
	ne->mac = mac;
	ne->next = 0;
	if (subn[hash2(ip)] != 0) {
		i = subn[hash2(ip)];

		while (i->next != 0) i = i->next;
		i->next = ne;
	}
	else subn[hash2(ip)] = ne;
}

void subnet_get(ip_address ip, mac_address &mac) {
	subnet_entry *i;

	if (subn[hash2(ip)] != 0) {
		i = subn[hash2(ip)];

		while (i) {
			if (ip.byte1 == i->ip.byte1 && ip.byte2 == i->ip.byte2 && ip.byte3 == i->ip.byte3 && ip.byte4 == i->ip.byte4) {
				mac.byte1 = i->mac.byte1;
				mac.byte2 = i->mac.byte2;
				mac.byte3 = i->mac.byte3;
				mac.byte4 = i->mac.byte4;
				mac.byte5 = i->mac.byte5;
				mac.byte6 = i->mac.byte6;
				return;
			}
		}
	}
	memset(&mac, 0xff, sizeof(mac));
}

u_short find_port(ip_address ip, mac_address mac, u_short port, int assigned) {
	table_entry *i, *ne;

	if (port2ent[hash_function(ip, port)] != 0) {
		i = port2ent[hash_function(ip, port)];

		while (i) {
			if (ip.byte1 == i->hsip.byte1 && ip.byte2 == i->hsip.byte2 && ip.byte3 == i->hsip.byte3 && ip.byte4 == i->hsip.byte4 && port == i->lanport)
				return i->wanport;
			i = i->next;
		}
	}

	ne = new table_entry;
	ne->hsip = ip;
	ne->hsmac = mac;
	ne->lanport = port;
	if (assigned == -1) {
		ne->wanport = 20000 + opened;
		opened ++;
	}
	else
		ne->wanport = assigned;
	ne->next = 0;

	port2ent[ne->wanport] = ne;
	if (port2ent[hash_function(ip, port)] == 0) 
		port2ent[hash_function(ip, port)] = ne;
	else {
		i = port2ent[hash_function(ip, port)];
		while (i->next != 0) i = i->next;
		i->next = ne;
	}

	return ne->wanport;
}

void load_default_port() {
	FILE *cfg;
	int port, ip1, ip2, ip3, ip4, mac1, mac2, mac3, mac4, mac5, mac6, hp;

	printf("Reading preset port mapping ...\n");
	cfg = fopen("config", "r");
	while (fscanf(cfg, "%d\t%d.%d.%d.%d:%d\t%02x:%02x:%02x:%02x:%02x:%02x", &port, &ip1, &ip2, &ip3, &ip4, &hp, &mac1, &mac2, &mac3, &mac4, &mac5, &mac6) != EOF)	{
		printf("PORT %d -> %d.%d.%d.%d:%d\n", port, ip1, ip2, ip3, ip4, hp);
		ip_address ip;
		mac_address mac;
		u_short wp, lp;

		ip.byte1 = ip1;
		ip.byte2 = ip2;
		ip.byte3 = ip3;
		ip.byte4 = ip4;

		mac.byte1 = mac1;
		mac.byte2 = mac2;
		mac.byte3 = mac3;
		mac.byte4 = mac4;
		mac.byte5 = mac5;
		mac.byte6 = mac6;

		wp = port; lp = hp;

		find_port(ip, mac, lp, wp);
	}
	printf("\n");
}

int main() {
	int a, b, c, d;
	pcap_t *wanhandle, *lanhandle;
	ip_address wlip, llip, gwip, wlmask;
	u_long selfsn;
	mac_address gwmac, wlmac, llmac;
	pcap_pkthdr *header;
	address_group t_adr;
	const u_char *pkt_data;
	u_char *new_pack = new u_char[65536];
	ip_header *iph;
	int res;

	printf("                              Naive NAT Proxy\n                  Author : Yuan Tian   Date : DEC-17-2007\n\n");
	load_default_port();
	set_interface(&wanhandle, &lanhandle, &wlip, &wlmask, &llip, &wlmac, &llmac);
	printf("Please input gateway IP address : ");
	scanf("%d.%d.%d.%d", &a, &b, &c, &d);
	gwip.byte1 = a;
	gwip.byte2 = b;
	gwip.byte3 = c;
	gwip.byte4 = d;

	printf("\nGetting gateway MAC ...\n");
	get_gwmac(wanhandle, gwip, &gwmac, wlip, wlmac);

	printf("\nWAN port IP : %d.%d.%d.%d\n",
		wlip.byte1,
		wlip.byte2,
		wlip.byte3,
		wlip.byte4
	);

	printf("WAN port Subnet Mask : %d.%d.%d.%d\n",
		wlmask.byte1,
		wlmask.byte2,
		wlmask.byte3,
		wlmask.byte4
	);

	selfsn = (*(u_long *)(&wlip)) & (*(u_long *)(&wlmask));

	printf("WAN port MAC :  %02x:%02x:%02x:%02x:%02x:%02x\n",
		wlmac.byte1,
		wlmac.byte2,
		wlmac.byte3,
		wlmac.byte4,
		wlmac.byte5,
		wlmac.byte6
	);

	printf("\nGateway IP : %d.%d.%d.%d\n",
		gwip.byte1,
		gwip.byte2,
		gwip.byte3,
		gwip.byte4
	);
	printf("Gateway MAC : %02x:%02x:%02x:%02x:%02x:%02x\n",
		gwmac.byte1,
		gwmac.byte2,
		gwmac.byte3,
		gwmac.byte4,
		gwmac.byte5,
		gwmac.byte6
	);

	printf("\nLAN port IP : %d.%d.%d.%d\n",
		llip.byte1,
		llip.byte2,
		llip.byte3,
		llip.byte4
	);
	printf("LAN port MAC :  %02x:%02x:%02x:%02x:%02x:%02x\n",
		llmac.byte1,
		llmac.byte2,
		llmac.byte3,
		llmac.byte4,
		llmac.byte5,
		llmac.byte6
	);

	while (1) {
		/* outward traffic */
		res = pcap_next_ex(lanhandle, &header, &pkt_data);
		if (res > 0  && *(pkt_data) == llmac.byte1 && *(pkt_data + 1) == llmac.byte2 && *(pkt_data + 2) == llmac.byte3 && *(pkt_data + 3) == llmac.byte4 && *(pkt_data + 4) == llmac.byte5 && *(pkt_data + 5) == llmac.byte6) {
			/* is IP packet */
			if (ntohs(*((u_short *)(pkt_data + 12))) == 0x0800) {
				iph = (ip_header *)(pkt_data + 14);

				/* is TCP */
				if (iph->proto == 6) {
					int tcp_start = 14 + 4 * (iph->ver_ihl & 0x0F);

					t_adr.isrc = wlip;
					t_adr.idest = iph->daddr;
					t_adr.msrc = wlmac;

					if ((*(u_long *)(&iph->daddr) & *(u_long *)(&wlmask)) == selfsn) {
						mac_address snm;

						subnet_get(iph->daddr, snm);
						t_adr.mdest = snm;
					}
					else
						t_adr.mdest = gwmac;

					t_adr.psrc = find_port(iph->saddr, *(mac_address *)(pkt_data + 6), ntohs(*(u_short *)(pkt_data + tcp_start)), -1);
					t_adr.pdest = ntohs(*(u_short *)(pkt_data + tcp_start + 2));

					regroup_tcp(pkt_data, &new_pack, header->len, t_adr);

					if (pcap_sendpacket(wanhandle, new_pack, header->len) != 0) {
						printf("error sending packet!");
						exit(0);
					}
				}

				/* is UDP */
				if (iph->proto == 17) {
					int udp_start = 14 + 4 * (iph->ver_ihl & 0x0F);

					t_adr.isrc = wlip;
					t_adr.idest = iph->daddr;
					t_adr.msrc = wlmac;
					if ((*(u_long *)(&iph->daddr) & *(u_long *)(&wlmask)) == selfsn) {
						mac_address snm;

						subnet_get(iph->daddr, snm);
						t_adr.mdest = snm;
					}
					else
						t_adr.mdest = gwmac;

					t_adr.psrc = find_port(iph->saddr, *(mac_address *)(pkt_data + 6), ntohs(*(u_short *)(pkt_data + udp_start)), -1);
					t_adr.pdest = ntohs(*(u_short *)(pkt_data + udp_start + 2));

					regroup_udp(pkt_data, &new_pack, header->len, t_adr);

					if (pcap_sendpacket(wanhandle, new_pack, header->len) != 0) {
						printf("error sending packet!");
						exit(0);
					}
				}
			}
		}

		/* inward traffic */
		res = pcap_next_ex(wanhandle, &header, &pkt_data);
		if (res > 0  && *(pkt_data) == wlmac.byte1 && *(pkt_data + 1) == wlmac.byte2 && *(pkt_data + 2) == wlmac.byte3 && *(pkt_data + 3) == wlmac.byte4 && *(pkt_data + 4) == wlmac.byte5 && *(pkt_data + 5) == wlmac.byte6) {

			if (ntohs(*((u_short *)(pkt_data + 12))) == 0x0800) {
				iph = (ip_header *)(pkt_data + 14);

				if ((*(u_long *)(&iph->saddr) & *(u_long *)(&wlmask)) == selfsn) {
					subnet_put(iph->saddr, *(mac_address *)(pkt_data + 6));
				}

				/* is TCP */
				if (iph->proto == 6) {
					int tcp_start = 14 + 4 * (iph->ver_ihl & 0x0F);

					if (port2ent[ntohs(*(u_short *)(pkt_data + tcp_start + 2))] != 0) {

						t_adr.isrc = iph->saddr;
						t_adr.idest = port2ent[ntohs(*(u_short *)(pkt_data + tcp_start + 2))]->hsip;
						t_adr.pdest = port2ent[ntohs(*(u_short *)(pkt_data + tcp_start + 2))]->lanport;
						t_adr.psrc = ntohs(*(u_short *)(pkt_data + tcp_start));
						t_adr.mdest = port2ent[ntohs(*(u_short *)(pkt_data + tcp_start + 2))]->hsmac;
						t_adr.msrc = llmac;
						regroup_tcp(pkt_data, &new_pack, header->len, t_adr);
						if (pcap_sendpacket(lanhandle, new_pack, header->len) != 0) {
							printf("error sending packet!");
							exit(0);
						}

					}
				}

				/* is UDP */
				if (iph->proto == 17) {
					int udp_start = 14 + 4 * (iph->ver_ihl & 0x0F);

					if (port2ent[ntohs(*(u_short *)(pkt_data + udp_start + 2))] != 0) {
						
						t_adr.isrc = iph->saddr;
						t_adr.idest = port2ent[ntohs(*(u_short *)(pkt_data + udp_start + 2))]->hsip;
						t_adr.pdest = port2ent[ntohs(*(u_short *)(pkt_data + udp_start + 2))]->lanport;
						t_adr.psrc = ntohs(*(u_short *)(pkt_data + udp_start));
						t_adr.mdest = port2ent[ntohs(*(u_short *)(pkt_data + udp_start + 2))]->hsmac;
						t_adr.msrc = llmac;
						regroup_udp(pkt_data, &new_pack, header->len, t_adr);
						if (pcap_sendpacket(lanhandle, new_pack, header->len) != 0) {
							printf("error sending packet!");
							exit(0);
						}

					}
				}

			}
		}
	}
}